from tqdm import tqdm
from yan import CvFo
import torch
import pandas as pd

voc = pd.read_pickle("voc_data.pandas_pickle")
net = CvFo(len(voc), 512, 8)
net.load_state_dict(torch.load("model_10_loss_3.76836.pth", map_location=torch.device('cpu')))
net.eval()
loss_func = torch.nn.CrossEntropyLoss()
data_set = pd.read_pickle("train_data.pandas_pickle")
batch_size = 1
bar = tqdm(range(0, len(data_set), batch_size))

for i in bar:
    j = i + batch_size
    one_data = data_set[i:j]
    print("".join([voc[i] for i in one_data[0][-150:-160]]))
    two_data = torch.Tensor(one_data)[:, :-150].int()
    for i in range(0, 10):
        out = net(two_data)
        token_id = torch.argmax(out[:, -1:, :], -1).item()
        print(voc[token_id],end="")
        two_data = torch.concat([two_data, torch.Tensor([[token_id]]).int()], -1)
